import os

from PIL import Image
from torch.utils.data import Dataset
from torchvision.transforms import transforms

train_real_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    # transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])

train_anime_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    # transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
train_gray_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.Grayscale(3),
    # transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])

val_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])


class AnimeDataset(Dataset):
    def __init__(self):
        super(AnimeDataset, self).__init__()
        root = "/home/pc/dataset/cartoon/"
        self.real_photo = self.getimages(root + "train_photo")
        kind = "hayao"
        self.anime_smooth = self.getimages(root+   "smooth")
        self.anime_style = self.getimages(root +  "style")

    def __len__(self):
        return len(self.anime_style)

    def __getitem__(self, index):
        n = len(self.real_photo)
        index = index % n
        photo_image = Image.open(self.real_photo[index]).convert("RGB")
        cartoon_smooth = Image.open(self.anime_smooth[index]).convert("RGB")
        cartoon = Image.open(self.anime_style[index]).convert("RGB")

        #                 photo                       cartoon          x grayscale animation training data    y remove edges of the animation images
        return train_real_transform(photo_image), train_real_transform(cartoon), train_gray_transform(cartoon), train_gray_transform(cartoon_smooth)

    def getimages(self, data_path):
        data = []
        for image in os.listdir(data_path):
            data.append(data_path + "/" + image)
        return data


class TestAnimeDataset(Dataset):
    def __init__(self):
        super(TestAnimeDataset, self).__init__()
        root = "F:/cvpr_dataset/AnimeGan/test/test_photo/"
        self.photo = self.getimages(root)

    def __len__(self):
        return len(self.photo)

    def __getitem__(self, index):
        n = len(self.photo)
        index = index % n
        photo_image = Image.open(self.photo[index]).convert("RGB")
        return val_transform(photo_image)

    def getimages(self, data_path):
        data = []
        for image in os.listdir(data_path):
            data.append(data_path + image)
        return data